import torch
import cv2
import numpy as np
import matplotlib

matplotlib.use('Agg')  # 设置非GUI后端
from matplotlib import pyplot as plt
from torchvision import transforms
from PIL import Image
from models import get_modified_resnet18


def preprocess_image(img_path):
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
    ])
    img = Image.open(img_path).convert('RGB')
    return img, preprocess(img).unsqueeze(0)


class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.activations = None
        self.gradients = None
        self.hooks = []
        self._register_hooks()

    def _register_hooks(self):
        def forward_hook(module, input, output):
            self.activations = output.detach()

        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0].detach()

        self.hooks.append(self.target_layer.register_forward_hook(forward_hook))
        self.hooks.append(self.target_layer.register_backward_hook(backward_hook))

    def __call__(self, input_tensor):
        self.model.eval()
        output = self.model(input_tensor)
        pred_class = output.argmax().item()
        self.model.zero_grad()
        output[0, pred_class].backward()

        pooled_grads = torch.mean(self.gradients, dim=[0, 2, 3])
        cam = torch.zeros(self.activations.shape[2:], device=input_tensor.device)
        for i in range(self.activations.shape[1]):
            cam += pooled_grads[i] * self.activations[0, i, :, :]

        cam = torch.relu(cam).cpu().numpy()
        cam = cv2.resize(cam, (224, 224))
        cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam) + 1e-8)
        return cam, pred_class

    def remove_hooks(self):
        for hook in self.hooks:
            hook.remove()


def save_cam_image(img, cam, save_path='gradcam.jpg'):
    img = np.array(img.resize((224, 224))) / 255.0
    heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
    blended = cv2.addWeighted(heatmap, 0.5, np.uint8(255 * img), 0.5, 0)
    cv2.imwrite(save_path, blended)
    return save_path


if __name__ == '__main__':
    import sys

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = get_modified_resnet18().to(device)
    model.load_state_dict(torch.load('best_model.pth', map_location=device))

    img_path = sys.argv[1] if len(sys.argv) > 1 else './data/0/image_0001.jpg'
    img, img_tensor = preprocess_image(img_path)
    img_tensor = img_tensor.to(device)

    target_layer = model.layer4[-1]
    gradcam = GradCAM(model, target_layer)
    cam, pred_class = gradcam(img_tensor)

    output_path = save_cam_image(img, cam)
    print(f"Grad-CAM saved to {output_path}")
    gradcam.remove_hooks()